import os
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate  # 使用最新导入路径
from src.core.memory import RedisMemory
from src.core.knowledge_base import KnowledgeBase
from src.core.database import DatabaseManager, MySQLConfig

# 从环境变量获取配置
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379")
DB_CONFIG = MySQLConfig(
    host=os.getenv("DB_HOST", "localhost"),
    port=int(os.getenv("DB_PORT", "3306")),
    user=os.getenv("DB_USER", "root"),
    password=os.getenv("DB_PASSWORD", ""),
    database=os.getenv("DB_NAME", "financial_data")
)

# 定义LLM提示模板
ANALYSIS_PROMPT = PromptTemplate(
    input_variables=["query", "context", "data"],
    template="分析{query}的市场趋势，参考信息：{context}，历史数据：{data}"
)

RISK_PROMPT = PromptTemplate(
    input_variables=["query", "market_context", "risk_docs"],
    template="评估{query}的投资风险，市场分析：{market_context}，风险框架：{risk_docs}"
)

ADVICE_PROMPT = PromptTemplate(
    input_variables=["query", "market_analysis", "risk_assessment"],
    template="为{query}生成投资建议，市场分析：{market_analysis}，风险评估：{risk_assessment}"
)

class AnalystAgent:
    """负责市场趋势分析的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.knowledge_base = KnowledgeBase()
        self.db = DatabaseManager(config=DB_CONFIG)
        self.llm_chain = LLMChain(llm=llm, prompt=ANALYSIS_PROMPT)
    
    def analyze_market_trends(self, query: str) -> dict:
        """分析市场趋势"""
        # 1. 从知识库获取相关市场信息
        context_results = self.knowledge_base.search(query)
        context = "\n".join([res['document'] for res in context_results])
        
        # 2. 从数据库获取历史数据
        historical_data = self.db.execute_query(
            f"SELECT * FROM financial_data WHERE symbol = '{query}'"
        )
        
        # 3. 使用LLM分析趋势
        analysis = self.llm_chain.run(
            query=query,
            context=context,
            data=historical_data
        )
        
        # 4. 存储分析结果到记忆
        self.memory.store(f"market_analysis:{query}", analysis)
        
        return analysis

class RiskAgent:
    """负责风险评估的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.knowledge_base = KnowledgeBase()
        self.llm_chain = LLMChain(llm=llm, prompt=RISK_PROMPT)
    
    def assess_risk(self, analysis_result: dict) -> dict:
        """评估投资风险"""
        # 1. 获取市场分析结果
        query = analysis_result.get('query')
        market_context = self.memory.retrieve(f"market_analysis:{query}")
        
        # 2. 获取风险相关文档
        risk_results = self.knowledge_base.search("风险评估框架")
        risk_docs = "\n".join([res['document'] for res in risk_results])
        
        # 3. 使用LLM评估风险
        risk_assessment = self.llm_chain.run(
            query=query,
            market_context=market_context,
            risk_docs=risk_docs
        )
        
        # 4. 存储风险评估结果
        self.memory.store(f"risk_assessment:{query}", risk_assessment)
        
        return risk_assessment

class AdvisorAgent:
    """提供投资建议的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.llm_chain = LLMChain(llm=llm, prompt=ADVICE_PROMPT)
    
    def provide_advice(self, risk_assessment: dict) -> str:
        """生成投资建议"""
        # 1. 获取市场分析和风险评估
        query = risk_assessment.get('query')
        market_analysis = self.memory.retrieve(f"market_analysis:{query}")
        risk_assessment_result = self.memory.retrieve(f"risk_assessment:{query}")
        
        # 2. 使用LLM生成投资建议
        investment_advice = self.llm_chain.run(
            query=query,
            market_analysis=market_analysis,
            risk_assessment=risk_assessment_result
        )
        
        return investment_advice